package com.shujia.core

import com.alibaba.fastjson.{JSON, JSONObject}
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

import scala.collection.mutable

object DemoKeyByProcess {
  def main(args: Array[String]): Unit = {
    /**
      * 实时统计道路的平均车速
      *
      */
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    env.setParallelism(1)

    val linesDS: DataStream[String] = env.socketTextStream("master", 8888)

    //取出道路编号和车速
    val roadAndSpeedDS: DataStream[(Long, Double)] = linesDS.map(line => {
      val carJson: JSONObject = JSON.parseObject(line)
      val roadId: Long = carJson.getLong("road_id")
      val speed: Double = carJson.getDouble("speed")
      (roadId, speed)
    })

    val keyByDS: KeyedStream[(Long, Double), Long] = roadAndSpeedDS.keyBy(_._1)

    val avgSpeedDS: DataStream[(Long, Double)] = keyByDS.process(new MyKeyByProcessFunction)

    avgSpeedDS.print()
    env.execute()
  }

}

class MyKeyByProcessFunction extends KeyedProcessFunction[Long, (Long, Double), (Long, Double)] {
  //用于保存总的车速

  /**
    * 这样写不对：同一个并行度中所有的key会共用同一个变量
    *
    */
  //var sumSpeed = 0.0
  //用于保存车的数量
  //var count = 0

  val roadSpeedMap = new mutable.HashMap[Long, (Double, Long)]

  /**
    * processElement： 将数据一条一条传递给processElement，
    *
    * @param value : 一行数据
    * @param ctx   ： 上下文对象
    * @param out   ： 用于将数据发生到下游
    */
  override def processElement(value: (Long, Double),
                              ctx: KeyedProcessFunction[Long, (Long, Double), (Long, Double)]#Context,
                              out: Collector[(Long, Double)]): Unit = {
    val (roadId, speed) = value

    //从hashmap中获取这个道路总的车速和总的车辆
    var (sumSpeed: Double, count: Long) = roadSpeedMap.getOrElse(roadId, (0.0, 0L))

    //将这一次的车速累加到sumSpeed上
    sumSpeed += speed
    //累加乘车的数量
    count += 1

    //将最新的总的车速和总的车辆保存到hashmap中
    roadSpeedMap.put(roadId, (sumSpeed, count))

    //计算平均车速
    val avgSpeed: Double = sumSpeed / count
    //将数据发生到下游
    out.collect((roadId, avgSpeed))
    println(roadSpeedMap)
  }
}
